package com.aij.starter.service;

import java.io.File;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.locks.ReentrantLock;

/**
 * LLM虚拟服务类
 * @author aij
 */
public abstract class LlmService {

    public static LlmService instance = null;

    public interface LLMCallback {
    	void run(String message);
    }

    public enum MessageType
    {
        Run,
        Finished,
    }

    public class InMessage
    {
        public MessageType type;
        public String message;
        public LLMCallback callback;
    }
    ReentrantLock mutex = new ReentrantLock();
    LinkedBlockingQueue<InMessage> inMessages = new LinkedBlockingQueue<InMessage>();
    Thread thread = null;
    boolean isClosing = false;
    boolean isRunning = false;

    protected abstract void LlmInit();
    protected abstract void LlmRun(String word);
    protected abstract String LlmPop();
    protected abstract void LlmFree();
    
    public static void Instance(String model, String type)
    {
    	switch (type) {
    		case "ChatGlmApi":
	            instance = new ChatApiService(model);
	            break;
    		case "ChatGlmStream":
	            instance = new ChatStreamService(model);
	            break;
	    	case "ChatGlm":
	        	if (new File(model).exists() == false) {
	        		System.out.println("Model file not exist: " + model);
	        		break;
	        	}
	            instance = new ChatGlmService(model);
	            break;
	    	case "Llama":
	        default:
	        	if (new File(model).exists() == false) {
	        		System.out.println("Model file not exist: " + model);
	        		break;
	        	}
	            instance = new LlamaService(model);
	            break;
    	}
    }

    public void Initialize()
    {
        // Initialize LLM
        LlmInit();

        // Start thread for processing received LLM strings
        thread = new Thread(new LlmRunnable());
        thread.start();
    }
    
    private class LlmRunnable implements Runnable {
    	 
        public void run()
        {
            try {
                boolean waiting = false;
                LLMCallback callback = null;
                SetRunning(true);
                SetClosing(false);
                while (IsClosing() == false)
                {
                    if (waiting)
                    {
                    	String text = LlmPop();
                        if (text != null)
                        {
                        	//System.out.println("text:" + text);
                            if (callback != null)
                            {
                                callback.run(text);
                            }
                        }
                        if (text == null || text == "") {
                        	waiting = false;
                        }
                    }
                    else if (HasMessage())
                    {
                        InMessage msg = PopMessage();
                        switch (msg.type)
                        {
	                        case Run:
                        		LlmRun(msg.message);
	                            callback = msg.callback;
	                            waiting = true;
	                            break;
							default:
                                waiting = true;
                                callback = null;
								break;
                        }
                    }
                    else
                    {
                        waiting = true;
                        Thread.sleep(1);
                    }
                }
                isRunning = false;
            } catch (Exception e) {
                System.out.println(e.getMessage() + "\n" + e.getStackTrace());
            }
        }
    }

    // Add word to LLM queue
    public void Run(String word, LLMCallback callback)
    {
        if (IsClosing() == true || IsRunning() == false) return;
        if (word == null || word.length() == 0) return;

        InMessage im = new InMessage();
        im.type = MessageType.Run;
        im.message = word;
        im.callback = callback;
        QueueMessage(im);
    }

    public void QueueMessage(InMessage im)
    {
        mutex.lock();
        inMessages.add(im);
        mutex.unlock();
    }

    private boolean HasMessage()
    {
        boolean ret = false;
        mutex.lock();
        ret = inMessages.size() > 0;
        mutex.unlock();
        return ret;
    }

    private InMessage PopMessage()
    {
        InMessage im = null;
        mutex.lock();
        im = inMessages.poll();
        mutex.unlock();
        return im;
    }

    public void SetClosing(boolean val)
    {
        mutex.lock();
        isClosing = val;
        mutex.unlock();
    }

    public void SetRunning(boolean val)
    {
        mutex.lock();
        isRunning = val;
        mutex.unlock();
    }

    public boolean IsClosing()
    {
    	boolean val;
        mutex.lock();
        val = isClosing;
        mutex.unlock();
        return val;
    }

    public boolean IsRunning()
    {
    	boolean val;
        mutex.lock();
        val = isRunning;
        mutex.unlock();
        return val;
    }

    public void Destroy()
    {
        if (!IsClosing()) LlmFree();
        SetClosing(true);

        int wait = 2000;
        if (thread != null) {
        	try {
				thread.join(wait);
				thread = null;
			} catch (InterruptedException e) {
				System.out.println(e.getMessage() + "\n" + e.getStackTrace());
            	System.out.println("LlmService didn't shut down in time.");
			}
        }
    }

}
